import os
import numpy as np
import pickle
import logging
from tqdm import tqdm
from lazypredict.Supervised import LazyClassifier
from sklearn.model_selection import train_test_split
import joblib

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# 定义文件大小检查函数
def get_size(file):
    size_bytes = os.path.getsize(file)
    size_megabytes = size_bytes / (1024 * 1024)
    return size_megabytes

# 定义数据加载函数
def load_data(x_path, y_path):
    with open(x_path, 'rb') as f:
        X = pickle.load(f)
    with open(y_path, 'rb') as f:
        y = pickle.load(f)

    X = np.array(X)
    y = np.array(y)
    
    if len(X.shape) == 1:
        X = X.reshape(-1, 1)

    return X, y

# 主函数，训练并保存最佳模型
def train_and_save_best_model(train_folder, dest_folder):
    x_file_path = os.path.join(dest_folder, "X.pkl")
    y_file_path = os.path.join(dest_folder, "y.pkl")

    # Step 1: 加载已保存的X和y数据
    if os.path.exists(x_file_path) and os.path.exists(y_file_path):
        logging.info("加载现有的X和y数据")
        logging.info(f"X文件大小: {get_size(x_file_path):.2f}MB")
        logging.info(f"y文件大小: {get_size(y_file_path):.2f}MB")
        X, y = load_data(x_file_path, y_file_path)
    else:
        logging.error("X.pkl或y.pkl文件不存在，请检查路径或生成数据。")
        return

    # Step 2: 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    # Step 3: 使用 LazyClassifier 训练模型
    clf = LazyClassifier(verbose=0, ignore_warnings=True, custom_metric=None)
    models, predictions = clf.fit(X_train, X_test, y_train, y_test)

    # Step 4: 打印所有模型的性能指标
    print(models)
    # Step 5: 找到准确率最高的模型
    best_model_name = models.index[0]
    best_accuracy = models.iloc[0]["Accuracy"]
    logging.info(f'准确率最高的模型是: {best_model_name}, 准确率: {best_accuracy}')

    # Step 6: 保存最佳模型
    best_model = clf.models[best_model_name]
    model_path = os.path.join(dest_folder, 'best_model.pkl')
    joblib.dump(best_model, model_path)
    logging.info(f"最佳模型已保存到 '{model_path}'.")
    return best_model, model_path

# 确保在主函数中调用 train_and_save_best_model 函数
if __name__ == "__main__":
    train_folder = "D:\work\cat\data\train"  # 训练集文件夹路径
    dest_folder = "D:\work\TXprojrct6"  # 保存X和Y的文件夹路径
    train_and_save_best_model(train_folder, dest_folder)